# -*- coding: utf-8 -*-
'''
@Time    : 24-5-24 下午15:37
@Author  : zhangweihao
@File    : tasks.py
'''
import datetime
import logging
import sys
sys.path.append('/home/tutu/公司/finance_backend/glbx_backend')
from celery_task.main import celery,Celery
import joblib
from tqdm import tqdm
import pandas as pd
import numpy as np
import time
from celery import group
from configs.config import Configs
from api.v2.risk_retrieval.service import *
workflow_name = '银行高管敏感信息分析V1.2'

@celery.task()
def asyn_create_knowledge_base(name):
    flag, token = login_dify(email=Configs().DIFY_EMAIL, passwd=Configs().DIFY_PASSWD)
    if not flag:
        logging.error('登录失败')
        return
    dataset_id = create_knowledge_base_by_token(name, token)
    if dataset_id:
        # TODO 将知识库加入到workflow中
        add_knowledge_base2workflow_by_workflow_name(dataset_id, workflow_name, token)

def if_name_repeat(name):
    db = DBSql()
    sql = f"""
        SELECT COUNT(*) FROM risk_monitoring_list WHERE name = '{name}'
    """
    with db.query_pd():
        lis = pd.read_sql_query(sql, con=db.conn)['count'].tolist()
    return all(lis)
        

@celery.task()
def asyn_delete_knowledge_base(name_list):
    flag, token = login_dify(email=Configs().DIFY_EMAIL, passwd=Configs().DIFY_PASSWD)
    if not flag:
        logging.error('登录失败')
        return
    for name in set(name_list):
        if not if_name_repeat(name):
            delete_knowledge_base(name, token)

@celery.task()
def workflow_asy(api_key_workflow, api_key_dataset, name, vip_name):
    # file_paths = search_keyword_in_files(Configs().MARKED_STATIC_PATH, name)
    # new
    file_paths = get_keyword_in_files_by_searcher(name)
    file_paths = [i.replace('./marked_files/', './static/marked_files/') for i in file_paths]
    logging.info(file_paths)
    filtered_file_paths = check_file_isnt_handled(file_paths)
    logging.info(filtered_file_paths)
    if not filtered_file_paths:
        return [], [], {}
    
# 1.首先查找该名称自己的知识库id，如果不存在，创建一个
    flag, token = login_dify(email=Configs().DIFY_EMAIL, passwd=Configs().DIFY_PASSWD)
    if not flag:
        logging.error('登录失败')
        return [], [], {}
    
    # 检查该人名称的知识库是否创建
    dataset_id = get_knowledge_base_id(name, token)
    if not dataset_id:
        dataset_id = create_knowledge_base_by_token(name, token)
    
    # 检查该人名称的知识库是否在workflow中
    if dataset_id not in get_knowledge_base_ids_from_workflow_draft(workflow_name, token):
        add_knowledge_base2workflow_by_workflow_name(dataset_id, workflow_name, token)

    upload_all_articals(filtered_file_paths, dataset_id, api_key_dataset)
    ducument_info, res, res_origin = workflow(api_key_workflow, name)
    return ducument_info, res, res_origin

def get_risk_count(res):
    count = 0
    if isinstance(res, list):
        for item in res:
            if isinstance(item, dict):
                if item.get("desc"):
                    count += 1
                else:
                    count += get_risk_count(item.get("children"))
            elif isinstance(item, list):
                count += get_risk_count(item)
    return count

@celery.task()
def beat_get_monitor_person_risk_count():
    logging.warning('开始更新检测人员的风险数')
    db = DBSql()
    sql = f"""
        SELECT name FROM risk_monitoring_list WHERE is_enable = 1
    """
    with db.query_pd():
        name_set = sorted(set(pd.read_sql_query(sql, con=db.conn)['name'].tolist()))
    for name in name_set:
        document_id, res, _ = workflow_asy(Configs().api_key_workflow, Configs().api_key_dataset, name, None)
        count = get_risk_count(res)
        update_time = datetime.datetime.now()
        update_risk_count(name,count,update_time)

def update_risk_count(name, count, update_time):
    db = DBSql()
    sql = """
        UPDATE risk_monitoring_list SET risk_count = %s, update_time = %s WHERE name = %s AND is_enable = 1
    """
    with db.insert(sql, (count, update_time, name)):
        pass

@celery.task()
def async_delete_document_when_handle_privacy(name_list):
    logging.warning(f'处理文档{name_list}，开始异步删除任务...')
    flag, token = login_dify(email=Configs().DIFY_EMAIL, passwd=Configs().DIFY_PASSWD)
    if not flag:
        logging.error('登录失败')
        return 
    for name in name_list:
        delete_document_when_handle_privacy(name,token)

@celery.task()
def workflow_blocking_asy(api_key_report, json_, mode):
    res = workflow_blocking(api_key_report, json_, mode)
    return res.get('data',{}).get('outputs',{}).get('output','')

if __name__ == "__main__":
    pass
    
